import cv2
import numpy as np
import glob


# 设置寻找亚像素角点的参数
# 搜索窗口大小：11×11
winSize = (11, 11)
# 搜索限制：无限制，可全图搜索
zeroZone = (-1, -1)
# 搜索停止条件：循环次数达到30次或精度达到0.001
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)

# 棋盘格规格
w = 11
h = 8

# 表示出世界坐标系中的棋盘格点
objp = np.zeros((w * h, 3), np.float32)

# mgrid后的结果：
# 两个数组，第一个表示横坐标，第二个表示纵坐标
# 例如，w = 3， h = 2时
# x = [[0, 0],
#      [1, 1],
#      [2, 2]]
# y = [[0, 1],
#      [0, 1],
#      [0, 1]]
# 此时w对应横坐标，为了正确表示世界坐标，必须先转置再提取坐标
objp[:, :2] = np.mgrid[0:w, 0:h].T.reshape(-1, 2)

# 棋盘格尺寸为0.1mm
square_size = 0.1

# 储存棋盘格角点的世界坐标和图像坐标对
objpoints = []   # 世界坐标系中的三维点
imgpoints = []   # 图像平面的二维点
# 加载图像
images = glob.glob('./RGB_camera_calib_img/*.png')
print('Stage 1 completed')


i = 0
for fname in images:
    # 读取图片
    img = cv2.imread(fname)
    # 保存图像的高与宽
    h1, w1 = img.shape[:2]
    # 转换为灰度图像，提升角点检测的精确性与稳定性
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 寻找灰度图像中的所有角点
    flags = (cv2.CALIB_CB_NORMALIZE_IMAGE +
             cv2.CALIB_CB_ADAPTIVE_THRESH +
             cv2.CALIB_CB_FILTER_QUADS +
             cv2.CALIB_CB_LARGER)
    ret, corners = cv2.findChessboardCornersSB(gray, (w, h), flags)

    # 如果找到角点，就把它存储起来
    if ret:
        print('i: ', i)
        i += 1
        cv2.cornerSubPix(gray, corners, winSize, zeroZone, criteria)
        # 追加
        objpoints.append(objp)
        imgpoints.append(corners)

        # 在图像上表示角点
        cv2.drawChessboardCorners(img, (w, h), corners, ret)
        cv2.namedWindow('findCorners', cv2.WINDOW_NORMAL)
        cv2.resizeWindow('findCorners', 640, 480)
        cv2.imshow('findCorners', img)
        cv2.waitKey(200)

cv2.destroyAllWindows()

print('Stage 2 completed')
print('Calibrating now...')

# 进行标定
ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(objpoints, imgpoints, gray.shape[::-1], None, None)

print('-' * 10 + 'results' + '-' * 10)
print("ret:\n", ret)
print("mtx:\n", mtx)  # 内参数矩阵
print("dist:\n", dist)  # 畸变系数   distortion cofficients = (k_1,k_2,p_1,p_2,k_3)
print("rvecs:\n", rvecs)  # 旋转向量
print("tvecs:\n", tvecs)  # 平移向量


# # 读取图片
# img = cv2.imread(images[0])
# # 记录图像的高与宽
# h1, w1 = img.shape[:2]
# # 转换为灰度图像，提升角点检测的精确性与稳定性
# gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
#
# # 寻找灰度图像中的所有角点
# flags = (cv2.CALIB_CB_NORMALIZE_IMAGE +
#          cv2.CALIB_CB_ADAPTIVE_THRESH +
#          cv2.CALIB_CB_FILTER_QUADS +
#          cv2.CALIB_CB_LARGER)
# ret, corners = cv2.findChessboardCornersSB(gray, (19, 15), flags)
#
# print(ret)
#
# # 如果找到角点，就把它存储起来
# if ret:
#     cv2.cornerSubPix(gray, corners, winSize, zeroZone, criteria)
#     # 追加
#     objpoints.append(objp)
#     imgpoints.append(corners)
#
#     # 在图像上表示角点
#     cv2.drawChessboardCorners(img, (w, h), corners, ret)
#     cv2.namedWindow('findCorners', cv2.WINDOW_NORMAL)
#     cv2.resizeWindow('findCorners', 640, 480)
#     cv2.imshow('findCorners', img)
#     cv2.waitKey(0)